#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <apr_pools.h>
#include <apr_env.h>
#include <apr_file_io.h>
#include <apr_network_io.h>
#include <apr_portable.h>
#include <apr_strings.h>
#include <apr_tables.h>
#include "pmi.h"
#include "util.h"

#define PMII_MAXLINE 1024

#define SINGLETON_INIT_BUT_NO_PM 1
#define NORMAL_INIT_WITH_PM		 2
static int PMI_initialized = 0;

static int PMI_kvsname_len_max = 0;
static int PMI_key_len_max = 0;
static int PMI_value_len_max = 0;

/*
 * This is not used in our implementation and issues warnings
 * static int PMI_iter_next_idx = 0;
 */
static int PMI_debug = 0;

static int PMI_rank = 0;
static int PMI_size = 1;
static int PMI_spawned = 0;

static apr_pool_t *pmi_pool = NULL;
static apr_file_t *pm_file = NULL;
static apr_socket_t *pm_socket = NULL;
static apr_table_t *pm_cmd = NULL;
static apr_table_t *pm_singleton_kvs = NULL;

static apr_status_t pm_gets(char *str, int len)
{
	if (pm_file == NULL)
	{
		return apr_socket_gets(str, len, pm_socket);
	}
	else
	{
		return apr_file_gets(str, len, pm_file);
	}
}

static apr_status_t pm_puts(const char *str)
{
	if (pm_file == NULL)
	{
		return apr_socket_puts(str, pm_socket);
	}
	else
	{
		return apr_file_puts(str, pm_file);
	}
}

static void parse_response(char *line)
{
	char *p, *q, *last;
	
	/* clear out the key value table */
	
	apr_table_clear(pm_cmd);
	
	/* get rid of the newline if there is one */
	
	p = strchr(line, '\n');
	if (p != NULL) *p = '\0';
	
	/* iterate over the key=value tokens, split them and put them
	   in the table */
	
	p = apr_strtok(line, " ", &last);
	while (p != NULL)
	{
		q = strchr(p, '=');
		if (q != NULL)
		{
			*q = '\0';
			apr_table_set(pm_cmd, p, q + 1);
		}
		p = apr_strtok(NULL, " ", &last);
	}	
}

static void connect_to_pm(const char *hostname, int portnum)
{
 	apr_sockaddr_t *sa;
	
	apr_sockaddr_info_get(&sa, hostname, APR_UNSPEC, portnum, 0,
		pmi_pool);
	apr_socket_create(&pm_socket, APR_INET, SOCK_STREAM, APR_PROTO_TCP,
		pmi_pool);
	apr_socket_opt_set(pm_socket, APR_TCP_NODELAY, 1);
	apr_socket_connect(pm_socket, sa);
}

static int handshake_with_pm(int id)
{
	char buf[PMII_MAXLINE];
	const char *value;
	apr_status_t status;
	
	/* send initialization acknowledgement to process manager */
	
	snprintf(buf, PMII_MAXLINE, "cmd=initack pmiid=%d\n", id);
	status = pm_puts(buf);
	if (status != APR_SUCCESS) return -1;

	/* receive response from process manager; should be an ack
	   followed by set commands for size, rank and debug */
	
	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	parse_response(buf);
	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "initack") != 0) return -1;

	/* read size from process manager */
	
	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "set") != 0) return -1;
	value = apr_table_get(pm_cmd, "size");
	PMI_size = atoi(value);
	
	/* read rank from process manager */
	
	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "set") != 0) return -1;
	value = apr_table_get(pm_cmd, "rank");
	PMI_rank = atoi(value);
	
	/* read debug flag from process manager */
	
	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "set") != 0) return -1;
	value = apr_table_get(pm_cmd, "debug");
	PMI_debug = atoi(value);
	
	return 0;
}

static int send_init_to_pm()
{
	char buf[PMII_MAXLINE];
	const char *value;
	apr_status_t status;
	
	snprintf(buf, PMII_MAXLINE,
		"cmd=init pmi_version=%d pmi_subversion=%d\n", PMI_VERSION,
		PMI_SUBVERSION);
	status = pm_puts(buf);
	if (status != APR_SUCCESS) return -1;
	
	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	parse_response(buf);
	
	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "response_to_init") != 0) return -1;
	
	return 0;
}

static int get_kvs_limits()
{
	char buf[PMII_MAXLINE];
	const char *value;
	apr_status_t status;
	
	snprintf(buf, PMII_MAXLINE, "cmd=get_maxes\n");
	status = pm_puts(buf);
	if (status != APR_SUCCESS) return -1;

	status = pm_gets(buf, PMII_MAXLINE);
	if (status != APR_SUCCESS) return -1;
	parse_response(buf);

	/* cmd=maxes kvsname_max=nnn keylen_max=nnn vallen_max=nnn */

	value = apr_table_get(pm_cmd, "cmd");
	if (strcmp(value, "maxes") != 0) return -1;

	value = apr_table_get(pm_cmd, "kvsname_max");
	PMI_kvsname_len_max = atoi(value);

	value = apr_table_get(pm_cmd, "keylen_max");
	PMI_key_len_max = atoi(value);

	value = apr_table_get(pm_cmd, "vallen_max");
	PMI_value_len_max = atoi(value);

	return 0;
}

int PMI_Init(int *spawned)
{
	apr_status_t status;
	char *p;
	int rank_and_size_set = 0;
	
	/* allocate a memory pool */

	status = apr_pool_create(&pmi_pool, NULL);
	if (status != APR_SUCCESS) return -1;
	apr_pool_tag(pmi_pool, "pmi");
	
	/* allocate a table to hold key-values pairs from process manager
	   responses */
	
	pm_cmd = apr_table_make(pmi_pool, 2);
	
	/* get debug flag from environment */

	if (apr_env_get(&p, "PMI_DEBUG", pmi_pool) == APR_SUCCESS)
	{
		PMI_debug = atoi(p);
	}
	
	/* get process manager file descriptor from environment */

	if (apr_env_get(&p, "PMI_FD", pmi_pool) == APR_SUCCESS)
	{
		/* we already have connection to the process manager via a
		   file descriptior */

		apr_os_file_t fd = atoi(p);
		apr_os_file_put(&pm_file, &fd, O_RDWR, pmi_pool);
	}
	else if (apr_env_get(&p, "PMI_PORT", pmi_pool) == APR_SUCCESS)
	{
		/* connect to the indicated port (hostname:port) to get process
		   manager file descriptor */

		char *addr, *scope_id;
		apr_port_t port;
		status = apr_parse_addr_port(&addr, &scope_id, &port, p,
			pmi_pool);
		if ((status == APR_SUCCESS) && (addr != NULL))
		{
			connect_to_pm(addr, port);
		}
		if (pm_socket == NULL)
		{
			return -1;
		}
		
		/* if we have a PMI ID from the environment, handshake with the
		   process manager to get size and rank */

		if (apr_env_get(&p, "PMI_ID", pmi_pool) == APR_SUCCESS)
		{
			if (handshake_with_pm(atoi(p)) == 0)
			{
				rank_and_size_set =1;
			}
		}
	}
	else
	{
		/* singleton init; process not started with mpiexec */
		
		PMI_size = 1;
		PMI_rank = 0;
		PMI_initialized = SINGLETON_INIT_BUT_NO_PM;
		PMI_kvsname_len_max = 256;
		PMI_key_len_max = 256;
		PMI_value_len_max = 256;
		pm_singleton_kvs = apr_table_make(pmi_pool, 2);
		
		return 0;
	}

	/* if rank and size are not already set from handshaking, we
	   set them from the environment */
	
	if (!rank_and_size_set)
	{
		if (apr_env_get(&p, "PMI_SIZE", pmi_pool) == APR_SUCCESS)
		{
			PMI_size = atoi(p);
		}
		else
		{
			PMI_size = 1;
		}
		if (apr_env_get(&p, "PMI_RANK", pmi_pool) == APR_SUCCESS)
		{
			PMI_rank = atoi(p);
		}
		else
		{
			PMI_rank = 1;
		}
	}
	
	/* initialize our connection to the process manager */
	
	send_init_to_pm();
	
	/* get max string lengths from process manager */
	
	get_kvs_limits();
	
	/* find out whether we're spawned or not from the environment */
	
	if (apr_env_get(&p, "PMI_SPAWNED", pmi_pool) == APR_SUCCESS)
	{
		PMI_spawned = atoi(p);
	}
	else
	{
		PMI_spawned = 0;
	}
	*spawned = (PMI_spawned != 0);
	
	/* at this point, we have initialized normally with a PM */
	
	PMI_initialized = NORMAL_INIT_WITH_PM;

	return 0;
}

int PMI_Initialized()
{
	return PMI_initialized;
}

int PMI_Get_size(int *size)
{
	*size = (PMI_initialized) ? PMI_size : 1;
	return 0;
}

int PMI_Get_rank(int *rank)
{
	*rank = (PMI_initialized) ? PMI_rank : 0;
	return 0;
}

int PMI_Barrier()
{
	if (PMI_initialized > 1) /* ignore for singleton init with no PM */
	{
		char buf[PMII_MAXLINE];
		const char *value;
		apr_status_t status;
		
		status = pm_puts("cmd=barrier_in\n");
		if (status != APR_SUCCESS) return -1;
		status = pm_gets(buf, PMII_MAXLINE);
		if (status != APR_SUCCESS) return -1;
		parse_response(buf);
		value = apr_table_get(pm_cmd, "cmd");
		if (strcmp(value, "barrier_out") != 0)
		{
			printf("expecting cmd=barrier_out, got \"%s\"\n", value);
			return -1;
		}
		else
		{
			return 0;
		}
	}
	else
	{
		return 0;
	}
}

int PMI_Finalize()
{
	if (PMI_initialized > 1) /* ignore for singleton init with no PM */
	{
		char buf[PMII_MAXLINE];
		const char *value;
		apr_status_t status;
		
		status = pm_puts("cmd=finalize\n");
		if (status != APR_SUCCESS) return -1;
		status = pm_gets(buf, PMII_MAXLINE);
		if (status != APR_SUCCESS) return -1;
		parse_response(buf);
		
		value = apr_table_get(pm_cmd, "cmd");
		if (strcmp(value, "finalize_ack") != 0)
		{
			printf("expecting cmd=finalize_ack, got %s\n", value);
			return -1;
		}
		else
		{
			if (pm_socket != NULL)
			{
				apr_socket_shutdown(pm_socket, APR_SHUTDOWN_READWRITE);
			}
			if (pm_file != NULL)
			{
				apr_file_close(pm_file);
			}
		}
	}

	apr_pool_destroy(pmi_pool);
	return 0;
}

int PMI_KVS_Get_my_name(char *kvsname)
{
	if (PMI_initialized == SINGLETON_INIT_BUT_NO_PM)
	{
		/* return a dummy name */
		strncpy(kvsname, "singinit_kvs_0", PMI_kvsname_len_max);
		return 0;
	}
	else
	{
		char buf[PMII_MAXLINE];
		const char *value;
		apr_status_t status;
		
		status = pm_puts("cmd=get_my_kvsname\n");
		if (status != APR_SUCCESS) return -1;
		status = pm_gets(buf, PMII_MAXLINE);
		if (status != APR_SUCCESS) return -1;
		parse_response(buf);
		
		value = apr_table_get(pm_cmd, "cmd");
		if (strcmp(value, "my_kvsname") != 0)
		{
			printf("expected cmd=get_my_kvsname, got: %s\n", value);
			return -1;
		}
		else
		{
			value = apr_table_get(pm_cmd, "kvsname");
			strncpy(kvsname, value, PMI_kvsname_len_max);
			return 0;
		}
	}
}

int PMI_KVS_Get_name_length_max()
{
	return PMI_kvsname_len_max;
}

int PMI_KVS_Get_key_length_max()
{
	return PMI_key_len_max;
}

int PMI_KVS_Get_value_length_max()
{
	return PMI_value_len_max;
}

int PMI_KVS_Create(char *kvsname)
{
	return 0;
	
}

int PMI_KVS_Destroy(const char *kvsname)
{
	return 0;
	
}

int PMI_KVS_Put(const char *kvsname, const char *key, const char *value)
{
	/* cmd=put kvsname=%s key=%s value=%s\n */

	if (PMI_initialized > 1) /* ignore for singleton init with no PM */
	{
		char buf[PMII_MAXLINE];
		const char *v;
		apr_status_t status;
		
		snprintf(buf, PMII_MAXLINE,
			"cmd=put kvsname=%s key=%s value=%s\n", kvsname,
			key, value);
		status = pm_puts(buf);
		if (status != APR_SUCCESS) return -1;
		status = pm_gets(buf, PMII_MAXLINE);
		if (status != APR_SUCCESS) return -1;
		parse_response(buf);

		v = apr_table_get(pm_cmd, "cmd");
		if (strcmp(v, "put_result") != 0)
		{
			printf("expecting cmd=put_result, got %s\n", v);
			return -1;
		}
		else
		{
			return 0;
		}
	}
	else
	{
		apr_table_set(pm_singleton_kvs, key, value);
		return 0;
	}
}

int PMI_KVS_Commit(const char *kvsname)
{
	return 0;
	
}

int PMI_KVS_Get(const char *kvsname, const char *key, char *value)
{
	/* cmd=get kvsname=%s key=%s\n */

	if (PMI_initialized > 1) /* ignore for singleton init with no PM */
	{
		char buf[PMII_MAXLINE];
		const char *v;
		apr_status_t status;
		
		snprintf(buf, PMII_MAXLINE,
			"cmd=get kvsname=%s key=%s\n", kvsname, key);
		status = pm_puts(buf);
		status = pm_gets(buf, PMII_MAXLINE);
		parse_response(buf);
		
		v = apr_table_get(pm_cmd, "cmd");
		if (strcmp(v, "get_result") != 0)
		{
			printf("expecting cmd=get_result, got %s\n", v);
			return -1;
		}
		else
		{
			int rc;
			v = apr_table_get(pm_cmd, "rc");
			rc = atoi(v);
			if (rc == 0)
			{
				v = apr_table_get(pm_cmd, "value");
				strncpy(value, v, PMI_value_len_max);
				return 0;
			}
			else
			{
				return -1;
			}
		}
	}
	else
	{
		const char *v = apr_table_get(pm_singleton_kvs, key);
		strncpy(value, v, PMI_value_len_max);
		return 0;
	}
}

int PMI_KVS_iter_first(const char *kvsname, char *key, char *value)
{
	return 0;
	
}

int PMI_KVS_iter_next(const char *kvsname, char *key, char *value)
{
	return 0;
	
}

int PMI_Abort(int exit_code, const char error_msg[])
{
	printf("aborting job: %s\n", error_msg);
	exit(exit_code);
	return -1;
}
